{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NERDataset(Dataset):\n",
    "    def __init__(self, data_file):\n",
    "        self.data = self.load_data(data_file)\n",
    "    \n",
    "    def load_data(self, data_file):\n",
    "        Data = {}\n",
    "        with open(data_file, 'rt', encoding='utf-8') as f:\n",
    "            df = f.read().split('\\n\\n')\n",
    "            for idx, line in enumerate(df):\n",
    "                if not line:\n",
    "                    break\n",
    "                sentence, labels, char_labels = '', [], []  # 新增char_labels\n",
    "                items = line.split('\\n')\n",
    "                \n",
    "                # 构建实体标签\n",
    "                entity_spans = []\n",
    "                for i, item in enumerate(items):\n",
    "                    if not item.strip():\n",
    "                        continue\n",
    "                    char, tag = item.split(' ')\n",
    "                    sentence += char\n",
    "                    \n",
    "                    if tag.startswith('B'):\n",
    "                        entity_spans.append([i, i, char, tag[2:]])\n",
    "                        categories.add(tag[2:])\n",
    "                    elif tag.startswith('I'):\n",
    "                        entity_spans[-1][1] = i\n",
    "                        entity_spans[-1][2] += char\n",
    "                \n",
    "                # 生成字符级标签序列\n",
    "                char_labels = ['O'] * len(sentence)\n",
    "                for start, end, _, label_type in entity_spans:\n",
    "                    for pos in range(start, end + 1):\n",
    "                        if pos == start:\n",
    "                            char_labels[pos] = f'B-{label_type}'\n",
    "                        else:\n",
    "                            char_labels[pos] = f'I-{label_type}'\n",
    "                \n",
    "                Data[idx] = {\n",
    "                    'sentence': sentence,\n",
    "                    'labels': entity_spans,  # 实体跨度信息\n",
    "                    'char_labels': char_labels  # 字符级BIO标签\n",
    "                }\n",
    "        \n",
    "        print(f'数据集中包含的实体类型有：{categories}')\n",
    "        return Data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['O', 'O', 'O', 'O', 'O']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "['O'] * 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sentence-len: 5\n",
      "3 4 口苦 临床表现\n"
     ]
    }
   ],
   "source": [
    "labels = [[3, 4, '口苦', '临床表现']]\n",
    "sentence = '现头昏口苦'\n",
    "# {'sentence': '目的观察复方丁香开胃贴外敷神阙穴治疗慢性心功能不全伴功能性消化不良的临床疗效', 'labels': [[4, 10, '复方丁香开胃贴', '中医治疗'], [20, 32, '心功能不全伴功能性消化不良', '西医诊断']]}\n",
    "# sentence = '目的观察复方丁香开胃贴外敷神阙穴治疗慢性心功能不全伴功能性消化不良的临床疗效'\n",
    "# labels = [[4, 10, '复方丁香开胃贴', '中医治疗'], [20, 32, '心功能不全伴功能性消化不良', '西医诊断']]\n",
    "\n",
    "print(f\"sentence-len: {len(sentence)}\")\n",
    "# 生成字符级标签序列\n",
    "char_labels = ['O'] * len(sentence)\n",
    "for start, end, _, label_type in labels:\n",
    "  print(start, end, _, label_type)\n",
    "  for pos in range(start, end + 1):\n",
    "    # print(f\"pos: {pos}\")\n",
    "    if pos == start:\n",
    "      char_labels[pos] = f'B-{label_type}'\n",
    "    else:\n",
    "      char_labels[pos] = f'I-{label_type}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['O', 'O', 'O', 'B-临床表现', 'I-临床表现']\n"
     ]
    }
   ],
   "source": [
    "print(char_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# step5 定义数据处理函数\n",
    "def collote_fn(batch_samples):\n",
    "  batch_sentences = []\n",
    "  batch_labels = []\n",
    "  \n",
    "  for sample in batch_samples:\n",
    "    batch_sentences.append(sample['sentence'])\n",
    "    batch_labels.append(sample['labels'])\n",
    "  \n",
    "  # 在这里进行tokenize\n",
    "  batch_inputs = tokenizer(\n",
    "    batch_sentences, \n",
    "    padding=True, \n",
    "    truncation=True, \n",
    "    return_tensors=\"pt\"\n",
    "  )\n",
    "  \n",
    "  return {\n",
    "    'input_ids': batch_inputs['input_ids'],\n",
    "    'attention_mask': batch_inputs['attention_mask'], \n",
    "    'labels': batch_labels  # 注意：这里可能需要进一步处理labels model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "  }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# step5 定义数据处理函数\n",
    "# 我们需要通过 DataLoader 库来按 batch 加载数据，\n",
    "# 并且将文本以及标签都转换为模型可以接受的输入形式。\n",
    "\n",
    "# def collote_fn(batch):\n",
    "#     input_ids = []\n",
    "#     attention_mask = []\n",
    "#     labels = []\n",
    "#     for item in batch:\n",
    "#         input_ids.append(item['input_ids'])\n",
    "#         attention_mask.append(item['attention_mask'])\n",
    "#         labels.append(item['labels'])\n",
    "#     return {\n",
    "#         'input_ids': torch.tensor(input_ids, dtype=torch.long),\n",
    "#         'attention_mask': torch.tensor(attention_mask, dtype=torch.long),\n",
    "#         'labels': torch.tensor(labels, dtype=torch.long)\n",
    "#     }\n",
    "\n",
    "\n",
    "# train_dataloader = DataLoader(train_data, batch_size=4, shuffle=True, collate_fn=collote_fn)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tuning_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5b5cfca19401015264b4741f08a5ab630bb44e69b46d09de08cfb7d2d4656fc3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
